package org.example.service;

import com.zaxxer.hikari.HikariDataSource;
import lombok.extern.log4j.Log4j2;
import org.example.mapper.DataSourceConfigMapper;
import org.example.model.DataSourceConfig;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author by PengHongChao
 * @date 2025/7/7 10:28
 */
@Component
@Log4j2
public class DataSourceManager {
    private final List<DataSourceProvider> providers;
    private final Map<Long, DataSource> dataSourceCache = new ConcurrentHashMap<>();
    private final Map<Long, JdbcTemplate> jdbcTemplateCache = new ConcurrentHashMap<>();
    private final Map<Long, NamedParameterJdbcTemplate> namedJdbcTemplateCache = new ConcurrentHashMap<>();

    public DataSourceManager(List<DataSourceProvider> providers) {
        this.providers = providers;
    }

    @Resource
    private DataSourceConfigMapper dataSourceConfigMapper;

    @PostConstruct
    public void init() {
        // 应用启动时加载所有数据源
        List<DataSourceConfig> configs = dataSourceConfigMapper.selectAllEnabled();
        for (DataSourceConfig config : configs) {
            registerDataSource(config);
        }
    }

    // 注册数据源
    public void registerDataSource(DataSourceConfig config) {
        try {
            DataSourceProvider provider = providers.stream()
                    .filter(p -> p.getType().equalsIgnoreCase(config.getType()))
                    .findFirst()
                    .orElseThrow(() -> new IllegalArgumentException("Unsupported data source type: " + config.getType()));

            DataSource dataSource = provider.createDataSource(config);
            dataSourceCache.put(config.getId(), dataSource);
            jdbcTemplateCache.put(config.getId(), new JdbcTemplate(dataSource));
            namedJdbcTemplateCache.put(config.getId(), new NamedParameterJdbcTemplate(dataSource));
            log.info("成功注册数据源: ID={}, 名称={}, 类型={}",
                    config.getId(), config.getName(), config.getType());
        } catch (IllegalArgumentException e) {
            log.error("注册数据源失败: ID={}, 名称={}", config.getId(), config.getName(), e);
        }
    }

    // 移除数据源
    public void unregisterDataSource(Long id) {
        DataSource ds = dataSourceCache.remove(id);
        jdbcTemplateCache.remove(id);
        namedJdbcTemplateCache.remove(id);
        if (ds instanceof HikariDataSource) {
            ((HikariDataSource) ds).close();
        }
        // 添加其他连接池类型的支持
        else if (ds instanceof AutoCloseable) {
            try {
                ((AutoCloseable) ds).close();
            } catch (Exception e) {
                log.error("关闭数据源失败", e);
            }
        }
    }

    public JdbcTemplate getJdbcTemplate(Long id) {
        return Optional.ofNullable(jdbcTemplateCache.get(id))
                .orElseThrow(() -> new IllegalArgumentException("DataSource not found: " + id));
    }

    public NamedParameterJdbcTemplate getNamedJdbcTemplate(Long id) {
        return Optional.ofNullable(namedJdbcTemplateCache.get(id))
                .orElseThrow(() -> new IllegalArgumentException("DataSource not found: " + id));
    }

    @PreDestroy
    public void destroy() {
        dataSourceCache.values().forEach(ds -> {
            if (ds instanceof AutoCloseable) {
                try {
                    ((AutoCloseable) ds).close();
                } catch (Exception e) {
                    // 记录日志
                }
            }
        });
        dataSourceCache.clear();
        jdbcTemplateCache.clear();
        namedJdbcTemplateCache.clear();
    }
}
